# cython: profile=True
from numpy import linspace,mgrid,zeros_like, zeros, save
from numpy.linalg import inv, det,norm
from numpy import dot,zeros,array,ones,argmax,zeros_like,shape,isnan
#from math import sqrt,sin,cos,exp,pi
from scipy.optimize import fmin_tnc, fmin_l_bfgs_b
import cython,time, pdb
if not cython.compiled:
    from math import *
from v3dsolver.fluid_elements.fluid_elems import Saffman,VortexPanel
from v3dsolver.neighbour_finder import Bin,NeighbourFinder
from v3dsolver.vector cimport Vec,Vec_new
#from lpsolver import solver
from lpsolver cimport csolver, puLPsolver 

from random import choice
from mpi4py import MPI
from cython.operator cimport dereference as deref
import time
#from enthought.mayavi import mlab

def lstsq_err(f,Aeq_np,beq_np):
    """
    Returns least square error
    """
    err = dot(Aeq_np,f) - beq_np
    f = sum(err**2)
    g = dot(err.T, Aeq_np) + dot(Aeq_np.T,err)
    return f

cdef extern from "stdlib.h":
    void free(void* ptr)
    void* malloc(size_t size)
    void* realloc(void* ptr,size_t size)
    
cdef unitsphere_possible_nbrs(int n):
    """
    Returns a list of uniformly spaced points on a unit sphere centered 
    at the origin. 
        
    Any new neighbour would be placed at one of these lat-long positions.
        
    n: number of new points on each latitude (same as no. of such latitudes)
    """
    cdef list pts = [] 
    cdef int i,j
    cdef double x,y,z
    cdef double PI = 3.1415926535897931
    cdef double dtheta = 2.*PI/n
    cdef double theta,phi
    cdef Vec newpt
    for i in range(n):
        theta=dtheta*i
        z = cos(theta)
        if i==0 or i==(n-1):
            # just add one particle at the pole
            x = 0.0
            y = 0.0
            newpt = Vec_new(x,y,z)
            pts.append(newpt)
        else:
            for j in range(n):
                phi=dtheta*j
                x=sin(theta)*cos(phi)
                y=sin(theta)*sin(phi)
                newpt = Vec_new(x,y,z)
                pts.append(newpt)
    return pts

cdef list kernel_nbrs = unitsphere_possible_nbrs(20)

cdef class VRTSolver:
    """
    This class is responsible for solving VRT around an element and providing
    vorticities obtained by the neighbours of that element.
    """
    def __init__(self,base_elem,neighbours,bin,h_nu,n=8):
        """
        Initialization of VRTSolver class. 
        
        base_elem: element for which VRT is to be solved
        neighbours: dict of neighbours to base_elem (each of form {hash:elem}) 
        bin: bin element containing all elements
        n: number of new possible neighbours on each latitude (same as no. of 
          such latitudes)
        """
        self.base_elem = base_elem
        self.nbrs = neighbours
        self.bin = bin
        self.maxindex = len(self.bin.elems) - 1
        self.cellsize = self.bin.cellsize
        self.h_nu = h_nu
        self.R = self.cellsize/h_nu
        self.rnbr = sqrt(8)
        self.n = n
        self.possible_nbrs = self.get_possible_nbrs(n) 
        self.comm = MPI.COMM_WORLD
    
    def get_possible_nbrs(self,n):
        """
        Returns a list of uniformly spaced points on a sphere centered at base
        elem and of radius cellsize. 
        
        Any new neighbour would be placed at one of these lat-long positions, 
        depending on distance of this new neighbour from other neighbours.
        
        n: number of new points on each latitude (same as no. of such latitudes)
        """
        global kernel_nbrs
        cdef list pts = []
        cdef Vec base_elem_pos = self.base_elem.pos
        cdef Vec kernel_nbr 
        cdef Vec new_nbr
        for i in range(len(kernel_nbrs)):
            kernel_nbr = kernel_nbrs[i]
            new_nbr = kernel_nbr*(self.rnbr*self.h_nu) + base_elem_pos
            pts.append(new_nbr)
        return pts
    
    cdef add_neighbour(self):
        """
        Adds a neighbour particle to base_elem in case the number of neighbours
        it already has is less than the number of equations of VRT. 
        
        It tries to add the neighbour such that the sum of its distance from 
        all the other neighbours is maximum. This approach tends to fill up 
        holes in the vorticity distribution.
        """
        # sum of distance of possible neighbours from all other neighbours
        
        l = len(self.possible_nbrs)
        dist_possible_nbrs = zeros(l,dtype='float64')
         
        for j in range(len(self.nbrs)):
            hashnbr = self.nbrs[j]
            nbr_pos = hashnbr['pos']
            for i in range(l):
                possible_nbr = self.possible_nbrs[i]
                rvec = possible_nbr - nbr_pos
                r = rvec.length()
                r_normalized = r/(self.rnbr*self.h_nu)
                dist_possible_nbrs[i] += (r_normalized - 1./(0.0001 + r_normalized))
        argnbr = argmax(dist_possible_nbrs)
        # argnbr = choice(range(len(kernel_nbrs)))
        """
        avg_nbrpos = Vec()
        for i in range(len(self.nbrs)):
            avg_nbrpos += self.nbrs[i]['pos'] - self.base_elem.pos
        avg_nbrpos = avg_nbrpos/len(self.nbrs)
        
        if avg_nbrpos.length()==0.:
            new_pos = self.base_elem.pos - Vec(1,0,0)*sqrt(8)*self.h_nu
        else:
            new_pos = self.base_elem.pos - (avg_nbrpos/avg_nbrpos.length())*sqrt(8)*self.h_nu
        """
        new_pos = self.possible_nbrs[argnbr]
        new_index = self.maxindex + 1 # assigning all newly created elems same index (for now) 
        new_hashnbr = {'pos':new_pos,'rank':self.comm.rank,
                       'type':Saffman,'index':new_index}
        self.nbrs.append(new_hashnbr)

    cdef get_fractions(self,double* f,double* pmax_f,double* pmin_f,double* pfsum):
        """
        Returns the fractions of vorticity obtained by solving the Vorticity 
        Redistribution Technique equations.
        """
        cdef int neq = 10              # number of VRT equations
        cdef int nvar = len(self.nbrs) # number of variables = number of neighbours
        
        cdef int i,j

        cdef double** Aeq = <double**>malloc(neq*sizeof(double))
        for i in range(neq):
            Aeq[i] = <double*>malloc(nvar*sizeof(double))
        cdef double* beq = <double*>malloc(neq*sizeof(double))
        
        cdef Vec eps
        cdef double eps1
        cdef double eps2
        cdef double eps3
        cdef Vec base_elem_pos = self.base_elem.pos
        cdef double h_nu = self.h_nu
        cdef Vec nbr_pos

        # making equation in terms of f_ij
        for j in range(nvar):
            nbr_pos = self.nbrs[j]['pos']
            eps=(base_elem_pos - nbr_pos)/h_nu
            eps1 = eps.x
            eps2 = eps.y
            eps3 = eps.z
            Aeq[0][j] = 1.
            Aeq[1][j] = eps1
            Aeq[2][j] = eps2
            Aeq[3][j] = eps3
            Aeq[4][j] = eps1*eps2
            Aeq[5][j] = eps2*eps3
            Aeq[6][j] = eps1*eps3
            Aeq[7][j] = eps1*eps1
            Aeq[8][j] = eps2*eps2
            Aeq[9][j] = eps3*eps3
        
        beq[0] = 1.  
        beq[1] = 0.
        beq[2] = 0.
        beq[3] = 0.
        beq[4] = 0.
        beq[5] = 0.
        beq[6] = 0.   
        beq[7] = 2.
        beq[8] = 2.
        beq[9] = 2.
        
        cdef double sum_Aeq_i
        for i in range(neq):
            sum_Aeq_i = 0.
            for j in range(nvar):
                sum_Aeq_i += Aeq[i][j]
            beq[i] = beq[i] - 0.5*sum_Aeq_i
        
        
        ## PYTHON VERSION OF LPSOLVER
        #print 'using python'
        """
        Aeq_np = zeros((neq,nvar))
        beq_np = zeros(neq)
        for i in range(neq):
            for j in range(nvar):
                Aeq_np[i][j] = Aeq[i][j]
            beq_np[i] = beq[i]
        w = solver(Aeq_np,beq_np)
        """
        ## PYTHON VERSION OF LPSOLVER
        
        cdef double* w = <double*>malloc((nvar+1)*sizeof(double))
        csolver(Aeq,beq,w,neq,nvar)
        #puLPsolver(Aeq,beq,w,neq,nvar)
        
        for j in range(nvar+1):
            f[j] = w[j] + 0.5
        
        cdef double max_f = deref(pmax_f)
        cdef double min_f = deref(pmin_f)
        cdef double fsum = deref(pfsum)
        pmax_f[0] = 0.
        pmin_f[0] = 1.
        pfsum[0] = 0.
        for j in range(nvar):
            if f[j]>pmax_f[0]:
                pmax_f[0] = f[j]
            if f[j]<pmin_f[0]:
                pmin_f[0] = f[j]
            pfsum[0] += f[j]
        f[nvar] = pmax_f[0]
        
        for i in range(neq):
            free(Aeq[i])
        free(Aeq)
        free(beq)
        #free(w)
    
    def visualize(self):
        """
        Visualize
        """
        x = zeros(len(self.nbrs))
        y = zeros(len(self.nbrs))
        z = zeros(len(self.nbrs))
        s = zeros(len(self.nbrs))
        for i,nbr in enumerate(self.nbrs):
            x[i] = nbr['pos'][0]
            y[i] = nbr['pos'][1]
            z[i] = nbr['pos'][2]
            s[i] = 1 if nbr['index']>self.maxindex else 0
        from enthought.mayavi import mlab
        mlab.points3d(x,y,z,s,scale_mode='none',scale_factor=0.01,colormap='jet')
        mlab.show()
    
    def solve_vrt(self):
        """
        Solves the VRT and returns a list of neighbours in the form:
        [nbrhash,nbrpos,nbrvor]; where nbrvor is component of vorticity 
        obtained by diffusing current base_elem.
        
        countnew: stores the count of newly created elements
        """
        
        cdef int nvar = len(self.nbrs)
        cdef double min_f
        cdef double max_f
        cdef double fsum
        
        countnew = 0
        countfailed = 0
        
        while nvar<11:
            countnew += 1
            self.add_neighbour()
            nvar += 1
        
        cdef double* f = <double*>malloc((nvar+1)*sizeof(double))
        self.get_fractions(f,&max_f,&min_f,&fsum)
        
        while (max_f-1)>1e-6 or (0-min_f)>1e-6 or abs(fsum-1.0)>1e-6:
            print 'failed'
            free(f)
            countfailed += 1
            self.add_neighbour()
            nvar += 1
            f = <double*>malloc((nvar+1)*sizeof(double))
            self.get_fractions(f,&max_f,&min_f,&fsum)
        
        self.visualize()
        print 'yo'
        
        vrt_nbrs = []
        cdef Vec base_elem_vor = self.base_elem.vor
        cdef dict hashnbr,newnbr
        cdef int i
        for i in range(len(self.nbrs)): 
            hashnbr = self.nbrs[i]
            newnbr = {'pos':hashnbr['pos'],'rank':hashnbr['rank'],
                      'type':hashnbr['type'],'index':hashnbr['index']}
            newnbr.update({'vor':f[i]*base_elem_vor})
            vrt_nbrs.append(newnbr)
        free(f)
        return vrt_nbrs,countnew,countfailed
    
cdef bint check_sameside(Vec p1,Vec p2,Vec a,Vec b):
    """
    Checks if p1 and p2 lie on the same side of the line joining a and b.
    """
    cdef Vec tmp1 = b - a
    cdef Vec tmp2 = p1 - a
    cdef Vec tmp3 = p2 - a
    cdef Vec cp1 = tmp1.dot(tmp2)
    cdef Vec cp2 = tmp1.dot(tmp3)
    if cp1.dot(cp2) >=0:
        return True
    else:
        return False
        
cdef bint check_pt_in_triangle(Vec p,Vec a,Vec b,Vec c):
    """
    Checks if point p lies inside the triangle formed by vertices a, b and c.
    """ 
    if (check_sameside(p,a, b,c) and check_sameside(p,b, a,c)
        and check_sameside(p,c, a,b)):
        return True
    else:
        return False

cdef bint is_inside(Vec pt,panel):
    """
    Returns True if pos is above or below the panel. Above or below is 
    decided by panel unitnormal which is always outward.
    """
    if (pt-panel.centroid).dot(panel.unitnormal) > 0:
        return False
    else:
        return check_pt_in_triangle(pt,panel.vertices[0],
                                    panel.vertices[1],panel.vertices[2])

class DiffusionManager:
    """
    This class is responsible for diffusion of vorticity of all elements. It 
    uses Vorticity Redistribution Technique (VRT) to simulate diffusion.
    """
    def __init__(self,elems,nu,dt,R,merge_frac,mode='p',geometry=None):
        """
        Initialization of DiffusionManager object.
        
        elems: list of elements to be diffused (only local elems in case of 
               parallel implementation)
        nu: kinematic viscosity of the fluid
        dt: timestep for simulation
        R: radius for binning (keep around 5*sqrt(nu*dt))
        merge_frac: fraction of binning cellsize within which particles
                    will be merged
        geometry: Geometry object (keep None for unbounded flows)
        """
        self.nu = nu
        self.dt = dt
        self.R = R
        self.cellsize = self.R*sqrt(nu*dt)
        self.merge_frac = merge_frac
        self.mode = mode
        self.elems = elems
        if geometry is None:
            self.panels = []
        else:
            self.panels = geometry.panels
        self.comm = MPI.COMM_WORLD
        
    def get_preadded_elems(self,elems):
        """
        Creates a finer bin based on given elems which is used to add
        neighbours in that part of the bin which is relatively empty. 
        Returns the list of new elems added in this way.
        """
        preadded_positions = []
        cdef Bin finebin = Bin(elems,self.cellsize/2.)
        is_done = {}
        cutoff = 0.5
        for i in finebin.first_elem_indices.keys():
            for j in finebin.nbrlist[i]:
                if not is_done.has_key(j):
                    is_done[j] = 1 # mark cell j as done
                    nbrelems_count = 0
                    if finebin.nbrlist.has_key(j):
                        nbrlist = finebin.nbrlist[j]
                    else:
                        nbrlist = finebin.get_nbrs(j)
                    for k in nbrlist:
                        try:
                            nbrelems_count += finebin.elem_count[k]
                        except KeyError:
                            pass
                    expected_elemcount = nbrelems_count/27.
                    #print 'expected_elemcount:', expected_elemcount
                    try:
                        imbalance = expected_elemcount - finebin.elem_count[j]
                    except KeyError:
                        imbalance = expected_elemcount
                    n_new = int(round(imbalance/cutoff))
                    #print 'n_new:', n_new
                    if n_new>0:
                        for l in range(n_new):
                            # randomly add a particle
                            newpos = finebin.get_newelem_pos(j)
                            preadded_positions.append(newpos)
        
        blob_rad = elems[0].rad
        preadded_elems = []
        for i in range(len(preadded_positions)):
            newelem = Saffman(pos=preadded_positions[i],vor=Vec(),
                              vel=Vec(),rad=blob_rad)
            preadded_elems.append(newelem)
        return preadded_elems

    def dispose_nbrs(self,new_nbrs,comm=MPI.COMM_WORLD):
        """
        Dispose obtained neighbours to appropriate locations. 
        
        * Remote processor vorticities are sent back
        * Vorticities obtained from remote processor are updated 
        * Already existing elems update their vorticities according to local 
          diffusion
        * Newly created elems are added to the same processor
        """
        send_data = []
        for i in range(comm.size):
            send_data.append([])
        
        newelems = []
        newvors = [Vec() for i in range(len(self.elems))]
        maxindex = len(self.bin.elems) - 1
        blob_rad = self.elems[0].rad
        
        for hashnbr in new_nbrs:
            # if belongs to same processor
            if hashnbr['rank']==comm.rank:
                # if already existing
                if hashnbr['index']<=maxindex:
                    newvors[hashnbr['index']] += hashnbr['vor']
                # if newly created
                else:
                    newelem = Saffman(pos=hashnbr['pos'],vel=Vec(),
                                      vor=hashnbr['vor'],rad=blob_rad)
                    newelems.append(newelem)
                    
            # if belonging to remote processor
            else:
                send_data[hashnbr['rank']].append(hashnbr)
        
        # send non-local data
        reqlist = []
        for i in range(len(send_data)):
            if not i==comm.rank:
                reqlist += [comm.isend(send_data[i],dest=i,tag=7)]
        
        # recv non-local data
        for i in range(comm.size):
            if not i==comm.rank:
                recvd_data = comm.recv(source=i,tag=7)
                for remote_hashnbr in recvd_data:
                    newvors[remote_hashnbr['index']] += remote_hashnbr['vor']
        MPI.Prequest.Waitall(reqlist)
        
        # assign obtained vorticities to all elems
        for i,elem in enumerate(self.elems):
            elem.vor = newvors[i]
        
        # append newly created particles to self.elems list
        self.elems += newelems

    def get_diffusable_elems(self,elems):
        """
        Returns elements which can be diffused. 
        
        Warning: Right now, only implemented for flow around a unit 
                 sphere centered at origin.
        """
        diffusable_elems = []
        for elem in elems:
            if elem.pos.length() >= (0.9398 + self.cellsize):
                diffusable_elems += [elem]
        #return diffusable_elems
        return elems
        
    def preprocessing(self):
        """
        Sets up the bin and neighbour finder for diffusion
        """
        self.bin = Bin(self.elems+self.panels,self.cellsize)
        self.nbrfinder = NeighbourFinder(self.bin,self.nu,self.dt,self.R)
        self.get_nbrs = {'d': self.nbrfinder.get_neighbours_direct,
                         's': self.nbrfinder.get_neighbours_serial,
                         'p': self.nbrfinder.get_neighbours_parallel}
        self.nbrs = {} # dict containing elem vs list of VRT generated nbrs in hash form
    
    def diffuse(self):
        """
        Performs diffusion using VRT. Returns the list of elems obtained 
        after diffusion.
        """
        h_nu = sqrt(self.nu*self.dt)
        new_nbrs = []
        
        totalcountnew = 0
        totalcountfailed = 0
        
        for i,elem in enumerate(self.elems):
            print i, elem.vor,
            nbrs = self.get_nbrs[self.mode](elem)
            #print i, '\t', len(nbrs)
            vrtsolver = VRTSolver(elem,nbrs,self.bin,h_nu)
            vrt_nbrs,countnew,countfailed = vrtsolver.solve_vrt()
            #print '\t', countnew, countfailed
            totalcountnew += countnew
            totalcountfailed += countfailed
            nbrvor = Vec()
            for nbr in vrt_nbrs:
                nbrvor += nbr['vor']
            print nbrvor
            try:
                assert (elem.vor-nbrvor).length()<=1e-6
            except AssertionError:
                print 'elemfailed:', (i, elem.pos, elem.vor)
                print 'vrt_nbrsfailed_index:', [nbr['index'] for nbr in vrt_nbrs]
                print '\nvrt_nbrsfailed_pos:', [(nbr['pos']-elem.pos)/self.cellsize for nbr in vrt_nbrs]
                print '\nvrt_nbrsfailed_vor:', [nbr['vor'] for nbr in vrt_nbrs]
                pdb.set_trace()
            new_nbrs += vrt_nbrs
            self.nbrs.update({elem:vrt_nbrs})
            
        
        # dispose obtained vorticities to appropriate particles
        self.dispose_nbrs(new_nbrs)
        print 'total count new:', totalcountnew
        print 'total count failed:', totalcountfailed
        
    def reflect_from_bdry(self):
        """
        Uses the neighbour finder to obtain blobs close to the boundary 
        panels and reflects those going inside the boundary.
        """
        count = 0
        for panel in self.panels:
            nbrs = self.get_nbrs[self.mode](panel)
            for hashnbr in nbrs:
                if hashnbr['type']==Saffman:
                    nbr = self.elems[hashnbr['index']]
                    if is_inside(nbr.pos,panel):
                        # normal component of blob velocity
                        velnormal = nbr.vel.dot(panel.unitnormal) 
                        
                        # normal component of panel pos w.r.t. blob
                        posnormal = (nbr.pos-panel.centroid).dot(panel.unitnormal)
                        
                        delta = max(velnormal*self.dt,posnormal)
                        nbr.pos = nbr.pos + 2*delta*panel.unitnormal
                        nbr.vel = nbr.vel - 2*velnormal
                        count += 1
        print 'reflect count:', count
                        
    def checkmerge(self,elem,nbr):
        """
        Checks if nbr can be merged to elem based on following conditions:
        
        1. nbr belongs to same processor as elem
        2. nbr is a dict based hashed representation of neighbor elem and
           the real neighbour elem is of type 'Saffman'
        3. nbr is within r_merge of elem
        """
        c1 = nbr['rank']==self.comm.rank
        c2 = nbr['type']==Saffman
        c3 = (elem.pos - nbr['pos']).length() <= sqrt(6)*sqrt(self.nu*self.dt)
        if (not c1) or (not c2) or (not c3):
            return False
        return True
    
    def merge(self):
        """
        Traverses in the obtained neighbour list for each particle to 
        check if they are close enough to be merged, and then merges 
        them.
        """
        # store which blob will be merged to which blob
        merged = zeros_like(self.elems) # 0 if elem not merged, 1 otherwise
        new_elems = []
        for i,elem in enumerate(self.elems):
            if self.nbrs.has_key(elem):
                nbrs = self.nbrs[elem]
            
                for nbr in nbrs:
                    if self.checkmerge(elem,nbr):
                        i_nbr = nbr['index']
                        if (merged[i]==0 and not i==i_nbr): 
                            merged[i] = 1
                            merged[i_nbr] = 1
                            elem1 = elem
                            elem2 = self.elems[i_nbr]
                            elem1vor = elem1.vor.length()
                            elem2vor = elem2.vor.length()
                            
                            newpos = ((elem1.pos*elem1vor + elem2.pos*elem2vor)/
                                      (elem1vor + elem2vor))
                            newvor = elem1.vor + elem2.vor
                            new_elems.append(Saffman(pos=newpos,vor=newvor,vel=Vec(),
                                                     rad=elem1.rad))
        
        # delete blobs which are merged
        for i in range(len(merged))[::-1]: # traversing in reverse order so as not to mess with the order
            if merged[i]==1:
                del self.elems[i]
        self.elems += new_elems
